# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : spark_ai_python.py
# @Author: dongguangwen
# @Date  : 2024-06-23 19:27

from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

#星火认知大模型Spark3.5 Max的URL值，其他版本大模型URL值请前往文档（https://www.xfyun.cn/doc/spark/Web.html）查看
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v1.1/chat'
# SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.5/chat'
#星火认知大模型调用秘钥信息，请前往讯飞开放平台控制台（https://console.xfyun.cn/services/bm35）查看
SPARKAI_APP_ID = '62e3e'
SPARKAI_API_SECRET = 'ZTUwYliMGIyMGVj'
SPARKAI_API_KEY = 'af78e3132b'

#星火认知大模型Spark3.5 Max的domain值，其他版本大模型domain值请前往文档（https://www.xfyun.cn/doc/spark/Web.html）查看
SPARKAI_DOMAIN = 'generalv3.5'

if __name__ == '__main__':
    spark = ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=False,
    )
    # messages = [ChatMessage(
    #     role="user",
    #     content='你好呀'
    # )]
    aa = '''永大纸胶带厂家直销电话
    永大纸胶带厂家直销店
    永大纸胶带厂家直销店地址
    永大纸胶带厂家直销店电话
    永大纸胶带厂家直销店在哪里
    永大纸胶带厂家直销价格
    永峰电工胶带厂家
    永峰电工胶带生产厂家
    永峰美纹胶带厂家
    永冠斑马警示胶带厂家
    永冠薄膜胶带厂家
    永冠保温胶带厂家
    永冠保温胶带厂家地址
    永冠保温胶带厂家地址查询'''
    messages = [ChatMessage(
        role="user",
        content=f'在输入的一些句子中，挑选出每句话中的品牌词。如”新锐胶带厂家“的产品词为”新锐“。输入的句子分别为：{aa}'
    )]
    handler = ChunkPrintHandler()
    a = spark.generate([messages], callbacks=[handler])
    print(a)
